#!/usr/bin/env python3
"""
简单测试跳层环境
"""
import numpy as np
import torch
from env.llm.model_loader import ModelLoader
from env.llm.skipable_state_llm import SkipableStateLLM
from env.env import SkipLayerEnv

def simple_env_test():
    """简单的环境测试"""
    print("开始简单环境测试...")
    
    # 加载模型
    model_loader = ModelLoader(cache_dir="./hf_cache/")
    model, tokenizer = model_loader.load_model("meta-llama/Llama-2-7b-chat-hf")
    
    # 创建模型实例
    full_llm = SkipableStateLLM(model, tokenizer)
    skip_llm = SkipableStateLLM(model, tokenizer)
    
    print(f"模型加载完成，层数: {full_llm.num_layers}")
    
    # 测试基本推理功能
    test_prompt = "Hello world"
    print(f"\n测试prompt: '{test_prompt}'")
    
    # 测试完整推理
    print("测试完整推理...")
    next_token, top_tokens, top_probs = full_llm.full_inference_with_state(test_prompt)
    print(f"完整推理结果: '{next_token}'")
    print(f"Top tokens: {top_tokens}")
    
    # 测试跳层推理
    print("\n测试跳层推理...")
    skip_llm.initialize(test_prompt)
    skip_llm.set_mask([False] * skip_llm.num_layers)  # 不跳过任何层
    skip_llm.inference(skip_llm.num_layers)
    skip_tokens, skip_probs = skip_llm.decode(k=5)
    print(f"跳层推理结果: '{skip_tokens[0]}'")
    print(f"一致性检查: {'✅' if next_token == skip_tokens[0] else '❌'}")
    
    # 创建环境
    print("\n创建环境...")
    env = SkipLayerEnv(
        full_model=full_llm,
        skip_model=skip_llm,
        dataset=[test_prompt],
        max_seq_length=5,
        top_k=3
    )
    
    # 测试环境
    print("测试环境...")
    state, info = env.reset()
    print(f"初始状态形状: {state.shape}")
    print(f"当前prompt: '{env.current_prompt}'")
    
    # 执行一步（不跳过任何层）
    action = np.array([0] * full_llm.num_layers)  # 不跳过任何层
    next_state, reward, done, truncated, info = env.step(action)
    
    print(f"执行动作后:")
    print(f"  生成的token: '{env.generated_sequence[-1] if env.generated_sequence else 'None'}'")
    print(f"  奖励: {reward:.4f}")
    print(f"  完成状态: {done}")
    print(f"  新状态形状: {next_state.shape}")
    
    # 执行另一步（跳过一些层）
    if not done:
        action = np.array([0, 1, 1] + [0] * (full_llm.num_layers - 3))  # 跳过前两层
        next_state, reward, done, truncated, info = env.step(action)
        
        print(f"\n跳层执行后:")
        print(f"  生成的token: '{env.generated_sequence[-1] if env.generated_sequence else 'None'}'")
        print(f"  奖励: {reward:.4f}")
        print(f"  完成状态: {done}")
    
    final_sequence = env.current_prompt + ''.join(env.generated_sequence)
    print(f"\n最终序列: '{final_sequence}'")
    print("测试完成!")

if __name__ == "__main__":
    simple_env_test()
